package org.hairball.viewer.prefuse;

import org.hairball.AnalysisType;
import org.hairball.ClassInclusionFilter;
import org.hairball.HairballGraphFactory;
import org.hairball.graph.EdgeVisitor;
import org.hairball.graph.HairballGraph;
import org.hairball.graph.NodeVisitor;
import prefuse.data.Edge;
import prefuse.data.Graph;
import prefuse.data.Node;
import prefuse.data.Schema;

import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class PrefuseGraphFactory {
    public Graph createGraph(String jarFileLocation) throws IOException {
        final Graph graph = new Graph();

        graph.getNodeTable().addColumns(LABEL_SCHEMA);
//        Node node1 = graph.addNode();
//        node1.setString(LABEL, "Thingy");
//
//        Node node2 = graph.addNode();
//        node2.setString(LABEL, "Thingy2");

        ClassInclusionFilter inclusionFilter = new ClassInclusionFilter() {

            public boolean shouldInclude(String fullyQualifiedClassName) {
                return true;
            }
        };
        Set<AnalysisType> analysisTypes = new HashSet<AnalysisType>();
        analysisTypes.add(AnalysisType.CONSTRUCTOR_INJECTION);
        analysisTypes.add(AnalysisType.SETTER_INJECTION);

        HairballGraph hairballGraph = new HairballGraphFactory().create(inclusionFilter, analysisTypes, jarFileLocation);
        final Map<String, Node> nodes = new HashMap<String, Node>();
        hairballGraph.visitNodes(new NodeVisitor() {
            public void visit(String nodeName, boolean connected) {
                Node node = graph.addNode();
                node.setString(LABEL, nodeName);
                nodes.put(nodeName, node);
            }
        });

        hairballGraph.visitEdges(new EdgeVisitor() {
            public void visit(String fromNodeName, String toNodeName, String edgeType) {
                Node fromNode = nodes.get(fromNodeName);
                Node toNode = nodes.get(toNodeName);
                Edge edge = graph.addEdge(fromNode, toNode);
                edge.setString(LABEL, edgeType);
            }
        });
        
        return graph;
    }

    public static final String LABEL = "label";
    /** Node table schema used for generated Graphs */
    public static final Schema LABEL_SCHEMA = new Schema();
    static {
        LABEL_SCHEMA.addColumn(LABEL, String.class, "");
    }
}
